# coding=utf-8
# 保存与加载 https://www.mindspore.cn/tutorials/zh-CN/master/beginner/save_load.html
import numpy as np
import mindspore
from mindspore import nn
from mindspore import Tensor


def network():
    model = nn.SequentialCell(
                nn.Flatten(),
                nn.Dense(28*28, 512),
                nn.ReLU(),
                nn.Dense(512, 512),
                nn.ReLU(),
                nn.Dense(512, 10))
    return model


# 保存和加载模型权重
model_dir = "../../model/model.ckpt"
model = network()
mindspore.save_checkpoint(model, model_dir)

param_dict = mindspore.load_checkpoint(model_dir)
param_not_load, _ = mindspore.load_param_into_net(model, param_dict)
# param_not_load是未被加载的参数列表，为空时代表所有参数均加载成功。
print(param_not_load)

# 保存和加载MindIR
# 除Checkpoint外，MindSpore提供了云侧（训练）和端侧（推理）统一的中间表示（Intermediate Representation，IR）。
# 可使用export接口直接将模型保存为MindIR。
model = network()
inputs = Tensor(np.ones([1, 1, 28, 28]).astype(np.float32))
mindspore.export(model, inputs, file_name="model", file_format="MINDIR")
# MindIR同时保存了Checkpoint和模型结构，因此需要定义输入Tensor来获取输入shape。
mindspore.set_context(mode=mindspore.GRAPH_MODE)

graph = mindspore.load("model.mindir")
model = nn.GraphCell(graph)
outputs = model(inputs)
print(outputs.shape)

